import numpy as np
from scipy.spatial import distance_matrix


def _calc_insert_cost(D, prv, nxt, ins):
    return (
            D[prv, ins]
            + D[ins, nxt]
            - D[prv, nxt]
    )


def run_insertion(loc):
    n = len(loc)
    D = distance_matrix(loc, loc)

    mask = np.zeros(n, dtype=bool)
    tour = []
    for i in range(n):
        feas = mask == 0
        # 函数输入一个矩阵，返回扁平化后矩阵中非零元素的位置（index）
        feas_ind = np.flatnonzero(mask == 0)
        if i == 0:
            a = D.max(1).argmax()  # Node with farthest distance to any other node
        else:
            a = feas_ind[D[np.ix_(feas, ~feas)].min(1).argmax()]  # node which has closest node in tour farthest

        mask[a] = True

        if len(tour) == 0:
            tour = [a]
        else:
            ind_insert = np.argmin(
                _calc_insert_cost(
                    D,
                    tour,
                    np.roll(tour, -1),
                    a
                )
            )
            tour.insert(ind_insert + 1, a)

    cost = D[tour, np.roll(tour, -1)].sum()
    return cost, tour


if __name__ == '__main__':
    city = np.array(
        [0.607122, 0.664447, 0.953593, 0.021519, 0.757626, 0.921024, 0.586376, 0.433565, 0.786837, 0.052959,
         0.016088, 0.581436, 0.496714, 0.633571, 0.227777, 0.971433, 0.665490, 0.074331, 0.383556, 0.104392])
    path = [1, 3, 8, 6, 10, 9, 5, 2, 4, 7]
    city = city.reshape(-1, 2)
    size = city.shape[0]
    cost, tour = run_insertion(city)
    print(tour, cost)
